import numpy as np

def loadDataSet(fileName):
    dataMat = []
    fr = open(fileName)
    for line in fr.readlines():
        curLine = line.strip().split('\t')
        fltLine = list(map(float, curLine)) #map all elements to float
        dataMat.append(fltLine)
    return dataMat

def binSplitDataSet(dataSet, feature, value):
    mat0 = dataSet[np.nonzero(dataSet[:, feature] > value)[0], :]
    mat1 = dataSet[np.nonzero(dataSet[:, feature] <= value)[0], :]
    return mat0, mat1

# 生成回归树叶节点，目标变量的均值
def regLeaf(dataSet): #returns the value used for each leaf
    return np.mean(dataSet[:, -1])

# 误差估计函数，使用均方差函数var()
# 因为要返回的是总方差，所以要用均方差乘以数据集中样本数
def regErr(dataSet):
    return np.var(dataSet[:,-1]) * np.shape(dataSet)[0]

# 对数据集求解线性回归模型，将数据集格式化成目标变量Y和自变量X
def linearSolve(dataSet):
    m, n = np.shape(dataSet)
    X = np.matrix(np.ones((m,n)))
    Y = np.matrix(np.ones((m,1)))
    X[:,1:n] = dataSet[:,0:n-1]
    Y = dataSet[:,-1]
    xTx = X.T * X
    if np.linalg.det(xTx) == 0.0:
        raise NameError('This matrix is singular, cannot do inverse, \n\
                        try increasing the second value of ops')
    ws = xTx.I * (X.T * Y)
    return ws, X, Y

# 模型树叶子节点保存的是线性回归模型的回归系数ws
def modelLeaf(dataSet):
    ws, X, Y = linearSolve(dataSet)
    return ws

# 模型树误差估计函数，先计算模型，然后预测，然后使用预测值和目标值计算平方误差和
def modelErr(dataSet):
    ws, X, Y = linearSolve(dataSet)
    yHat = X * ws
    return np.sum(np.power(Y-yHat, 2))


# tolS 是容许的误差下降值
# tolN 是切分的最少样本数
def chooseBestSplit(dataSet, leafType=regLeaf, errType=regErr, ops=(1,4)):
    tolS = ops[0]; tolN = ops[1]
    #if all the target variables are the same value: quit and return value
    #如果所有值相等则退出
    if len(set(dataSet[:, -1].T.tolist()[0])) == 1:
        #print('exit1')
        return None, leafType(dataSet)
    m, n = np.shape(dataSet)
    #the choice of the best feature is driven by Reduction in RSS error from mean
    S = errType(dataSet)
    bestS = np.inf; bestIndex = 0; bestValue = 0
    for featIndex in range(n-1):
        #print('featIndex', featIndex)
        for splitVal in set(dataSet[:, featIndex].T.tolist()[0]):            
            mat0, mat1 = binSplitDataSet(dataSet, featIndex, splitVal)
            #print(np.shape(mat0), np.shape(mat1), tolN)
            if (np.shape(mat0)[0] < tolN) or (np.shape(mat1)[0] < tolN):                            
                continue
            newS = errType(mat0) + errType(mat1)            
            if newS < bestS: #检查新切分是否能降低误差
                bestIndex = featIndex
                bestValue = splitVal
                bestS = newS
    #if the decrease (S-bestS) is less than a threshold don't do the split    
    #如果误差减少不大则退出
    if (S-bestS) < tolS:
        #print('exit2', S, bestS, S-bestS, tolS)
        return None, leafType(dataSet)
    mat0, mat1 = binSplitDataSet(dataSet, bestIndex, bestValue)
    #如果切分出的数据集很小则退出
    if (np.shape(mat0)[0] < tolN) or (np.shape(mat1)[0] < tolN):
        #print('exit3')
        return None, leafType(dataSet)    
    return bestIndex, bestValue

def createTree(dataSet, leafType=regLeaf, errType=regErr, ops=(1,4)):
    feat, val = chooseBestSplit(dataSet, leafType, errType, ops)
    if feat==None: return val #if the splitting hit a stop condition return val
    retTree = {}
    retTree['spInd'] = feat
    retTree['spVal'] = val
    lSet, rSet = binSplitDataSet(dataSet, feat, val)
    retTree['left'] = createTree(lSet, leafType, errType, ops)
    retTree['right'] = createTree(rSet, leafType, errType, ops)
    return retTree

def isTree(obj):
    return (type(obj).__name__=='dict')

# 求叶节点平均值，该操作会对树进行塌陷
def getMean(tree):
    print(tree)
    if isTree(tree['right']): tree['right'] = getMean(tree['right'])
    if isTree(tree['left']): tree['left'] = getMean(tree['left'])
    return (tree['left'] + tree['right'])/2.0

# 后剪枝
def prune(tree, testData):
    #if we have no test data collapse the tree
    if np.shape(testData)[0] == 0:
        return getMean(tree)
    
    if (isTree(tree['right']) or isTree(tree['left'])):
        lSet, rSet = binSplitDataSet(testData, tree['spInd'], tree['spVal'])
    if isTree(tree['left']):
        tree['left'] = prune(tree['left'], lSet)
    if isTree(tree['right']):
        tree['right'] = prune(tree['right'], rSet)
    #如果左右都是叶子，判断是否可以合并
    if not isTree(tree['left']) and not isTree(tree['right']):
        lSet, rSet = binSplitDataSet(testData, tree['spInd'], tree['spVal'])
        #不合并时的误差
        errorNoMerge = np.sum(np.power(lSet[:,-1] - tree['left'], 2)) + \
                        np.sum(np.power(rSet[:,-1] - tree['right'], 2))
        #合并时的误差
        treeMean = (tree['left']+tree['right'])/2.0
        errorMerge = np.sum(np.power(testData[:,-1] - treeMean, 2))
        if errorMerge < errorNoMerge:
            print("merging")
            return treeMean
        else: return tree
    else: return tree

def regTreeEval(model, inDat):
    return float(model)

def modelTreeEval(model, inDat):
    n = np.shape(inDat)[1]
    X = np.matrix(np.ones((1, n+1))) #增加第0列，全1
    X[:, 1:n+1]=inDat
    return float(X*model)

def treeForeCast(tree, inData, modelEval=regTreeEval):
    if not isTree(tree): return modelEval(tree, inData)
    if inData[tree['spInd']] > tree['spVal']:
        if isTree(tree['left']): return treeForeCast(tree['left'], inData, modelEval)
        else: return modelEval(tree['left'], inData)
    else:
        if isTree(tree['right']): return treeForeCast(tree['right'], inData, modelEval)
        else: return modelEval(tree['right'], inData)

def createForeCast(tree, testData, modelEval=regTreeEval):
    m = len(testData)
    yHat = np.matrix(np.zeros((m,1)))
    for i in range(m):
        yHat[i,0] = treeForeCast(tree, np.matrix(testData[i]), modelEval)
    return yHat